package org.basis.enhance.concurrent.executor;

import org.apache.commons.collections4.CollectionUtils;
import org.basis.enhance.concurrent.util.ListSplitUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

/**
 * 并发执行器
 *
 * @author Mr_wenpan@163.com 2020/2/4 2:27 下午
 */
public class ConcurrentExecutor<T> {

    private static final Logger LOG = LoggerFactory.getLogger(ConcurrentExecutor.class);

    private static final int DEFAULT_TASK_COUNT = 4;
    /**
     * 默认一个线程一次执行处理的最大数量
     */
    private static final int DEFAULT_BATCH_SIZE = 500;

    private static final String THREAD_FINISHED = "FINISHED";

    /**
     * 线程池中最大并行的任务数量
     */
    private final long maxTaskCount;
    /**
     * 每个任务执行的数据量大小
     */
    private final long batchSize;
    /**
     * 待处理的数据集合
     */
    private final List<T> processDataList;
    /**
     * 线程池
     */
    private final ThreadPoolExecutor taskExecutor;
    /**
     * 任务计数器，用于统计当前线程池里正在运行的线程数量计数
     */
    private final AtomicInteger taskCounter;

    public ConcurrentExecutor(int taskCount, List<T> processDataList, ThreadPoolExecutor taskExecutor) {
        this(taskCount, DEFAULT_BATCH_SIZE, processDataList, taskExecutor);
    }

    public ConcurrentExecutor(long maxTaskCount, long batchSize, List<T> processDataList, ThreadPoolExecutor taskExecutor) {
        this.maxTaskCount = maxTaskCount;
        this.batchSize = batchSize;
        this.processDataList = processDataList;
        this.taskExecutor = taskExecutor;
        taskCounter = new AtomicInteger(0);
    }

    /**
     * 并发执行任务
     *
     * @param consumer 待执行的任务逻辑
     * @author wenpan 2022/12/31 12:39 下午
     */
    public void invoke(Consumer<T> consumer) {
        // 如果待处理的数据为空，则直接返回
        if (CollectionUtils.isEmpty(processDataList)) {
            return;
        }

        // 如果指定了 maxTaskCount 最大为1，或者待处理的数据量小于等于每个线程一次最大处理的数据量，则直接由提交者线程处理
        if (maxTaskCount <= 1 || processDataList.size() <= batchSize) {
            doHandleTask(consumer, processDataList);
            return;
        }

        // 分发任务
        List<Future<String>> results = new ArrayList<>();
        // 将待处理的数据集合按照batchSize进行切分
        List<List<T>> subPendingLists = ListSplitUtil.splitList(processDataList, (int) batchSize);
        for (List<T> currentBatchProcess : subPendingLists) {
            // 将切分后的数据逐个提交到线程池，并将提交到线程池后返回的结果（Future<String>）添加到results中进行统一管理
            results.add(taskRunner(consumer, new ArrayList<>(currentBatchProcess)));
        }

        // 等待调度的线程执行结束
        for (Future<String> result : results) {
            try {
                result.get();
            } catch (Exception e) {
                LOG.error("ConcurrentExecutor invoke error.", e);
            }
        }
    }

    /**
     * 提交任务到线程池进行运行
     *
     * @param consumer   任务处理逻辑
     * @param processVar 待处理的任务数据
     * @return java.util.concurrent.Future<java.lang.String>
     * @author wenpan 2022/12/31 12:44 下午
     */
    private Future<String> taskRunner(Consumer<T> consumer, Collection<T> processVar) {
        // todo 这里一直循环设计不太优雅，后续优化为wait-notify模式
        while (true) {
            // 如果当前线程池中运行的线程数量小于线程池允许运行的最大任务数，则直接提交到线程池
            if (taskCounter.get() < maxTaskCount) {
                return taskExecutor.submit(() -> {
                    // 将当前线程池正在运行的任务数量++
                    taskCounter.incrementAndGet();
                    try {
                        // 任务处理逻辑
                        doHandleTask(consumer, processVar);
                    } catch (Exception e) {
                        LOG.error("ConcurrentExecutor taskRunner error.", e);
                    } finally {
                        taskCounter.decrementAndGet();
                    }
                    // 返回任务已经完成
                    return THREAD_FINISHED;
                });
            }
        }
    }

    private void doHandleTask(Consumer<T> consumer, Collection<T> processVar) {
        for (T processData : processVar) {
            consumer.accept(processData);
        }
    }
}